import tushare as ts
import sqlite3
import pandas as pd
import datetime
import math
import time

from comm.tsapis import Tsapi, logger

tb_nm = 'stock_trade'


# dbpath="/Users/hua/documents/TestData.db"

def getTag(x):
    tag = math.floor(x)
    if tag >= 8:
        return 8
    elif tag <= -8:
        return -8
    else:
        return tag


def tableCheck(conn):
    curs = conn.cursor()
    sql = f"SELECT name num FROM sqlite_master WHERE type='table' AND name='{tb_nm}'"
    curs.execute(sql)
    isexists = curs.fetchall()
    if len(isexists) > 0:
        curs.close()
    else:
        sql = f"""CREATE TABLE {tb_nm} (
         trade_date integer ,
         open REAL,
         high REAL,
         close REAL,
         pre_close REAL,
         low REAL,
         vol REAL,
         change REAL,
         pct_chg REAL,
         amount REAL,
         ma5 REAL,
         ma10 REAL,
         ma20 REAL,
         v_ma5 REAL,
         v_ma10 REAL,
         v_ma20 REAL,
         turnover REAL,
         ts_code TEXT,
         tag REAL
         )"""
        curs.execute(sql)
        curs.close()


def getLastDate(conn):
    curs = conn.cursor()
    sql = f"SELECT ts_code,ifnull(max(trade_date),20200101) maxdate FROM {tb_nm} group by ts_code"
    curs.execute(sql)
    dates = curs.fetchall()
    if len(dates) > 0:
        return dict(dates)
    else:
        return {'1000': 20200101}


def getTrade():
    hh = datetime.datetime.now().hour
    diff = 0 if hh >= 18 else -1
    yes_time = datetime.datetime.now() + datetime.timedelta(days=diff)
    enddate = int(yes_time.strftime("%Y%m%d"))
    ts = Tsapi()
    stocks = ts.get_stock_list()
    tableCheck(ts.conn.get_conn_refresh())
    mapdate = getLastDate(ts.conn.get_conn_refresh())
    i = 1
    total = len(stocks)
    stlist = list(stocks.ts_code)
    stlist.sort()
    stlist.append('sh')
    stlist.append('sz')
    logger.info("开始获取最近的交易数据，需要花点时间请等待完成")
    conn = ts.conn.get_conn_refresh()
    try:
        for code in stlist:
            if code in mapdate.keys():
                startdate = mapdate[code]
            else:
                startdate = 20220101
            i = i + 1
            logger.info(f"正在获取第{i},个,共：{total},-----{code} : from {startdate} to {enddate}")
            if enddate > startdate:
                d = ts.apis.query('daily', ts_code=code, start_date=startdate, end_date=enddate)
                # ts.get_hist_data(code,startdate,enddate)
                if d is not None and len(d) > 0:
                    # d['code'] = code
                    d['tag'] = d['pct_chg'].apply(getTag)
                    if 'turnover' not in d.columns:
                        d['turnover'] = math.nan
                    d.to_sql(tb_nm, conn, if_exists='append', index=False)
                    time.sleep(0.2)
    except Exception as e:
        logger.error(str(e))
        raise e
    finally:
        conn.close()
    logger.info("股票交易数据更新完成")


if __name__ == '__main__':
    getTrade()
